// Copyright (c) 2017 Franka Emika GmbH
// Use of this source code is governed by the Apache-2.0 license, see LICENSE
#pragma once

#include <array>
#include <cmath>

#include <Eigen/Core>
#include <franka/control_types.h>
#include <franka/duration.h>
#include <franka/robot.h>
#include <franka/robot_state.h>

void setDefaultBehavior(franka::Robot& robot);


class CubicTrajFunc{
  public:
    CubicTrajFunc(double vel_m, double acc_m, double distance)
        :vel_m_(vel_m), acc_m_(acc_m){
        sign_ = (distance < 0) ? -1 : 1;
        distance_ = sign_ * distance;

        dist_acc_ = 3 * std::pow(vel_m_, 2) / (4 * acc_m_);
        if(distance_ <= 2*dist_acc_){
            dist_acc_ = distance_ / 2;
            vel_m_ = std::sqrt(2 * acc_m_ * distance_ / 3);
        }
        t_acc_ = 3 * vel_m_ / (2 * acc_m_);
        t_uni_ = (distance_ - 2 * dist_acc_) / vel_m_;
        t_f_ = 2 * t_acc_ + t_uni_;
    }

    CubicTrajFunc(){};

    double operator()(double t){
        if (t < 0) {
            throw std::invalid_argument("time < 0");
        }
        double output;
        if(t < t_acc_){
            output = fun_(t);
        } else if(t < t_acc_ + t_uni_){
            output = dist_acc_ + vel_m_ * (t - t_acc_);
        } else if(t <= t_f_){
            output = distance_ - fun_(t_f_ - t);
        } else{
            output = distance_;
        }
        return sign_ * output;
    }

    double get_t_f(){
        return t_f_;
    }
    // 同步各个自由度的运动时间
    void set_t_f(double t_f_max){
        if(t_f_ > t_f_max){
            throw std::invalid_argument("t_f_ > t_f_max");
        } else if(t_f_ == t_f_max)
            return;
        t_f_ = t_f_max;
        vel_m_ = acc_m_ * (t_f_ - std::sqrt(std::pow(t_f_, 2) - 6*distance_/acc_m_)) / 3;
        dist_acc_ = 3 * std::pow(vel_m_, 2) / (4 * acc_m_);
        t_acc_ = 3 * vel_m_ / (2 * acc_m_);
        t_uni_ = (distance_ - 2 * dist_acc_) / vel_m_;
    }

  private:
    double vel_m_;  // 最大速度
    double acc_m_;  // 最大加速度
    double t_acc_;  // 加速阶段时长
    double t_uni_;  // 匀速阶段时长
    double t_f_;    // 总计时长
    double distance_;  // 运动距离
    double dist_acc_;  // 加速阶段运动距离
    double sign_;   // 运动方向正负
    inline double fun_(double t){
        return acc_m_ * std::pow(t, 3) * (2*t_acc_ - t) / (3 * std::pow(t_acc_, 2));
    }
};


class CosineTrajFunc{
  public:
    CosineTrajFunc(double vel_m, double acc_m, double distance){
        sign_ = (distance < 0) ? -1 : 1;
        distance_ = std::abs(distance);
        T_ = M_PI * distance_ / (2 * vel_m);
    };

    CosineTrajFunc(){};

    double operator()(double t){
        if(t < T_){
            return sign_ * distance_ / 2 * (1 - std::cos(M_PI / T_ * t));
        } else {
            return sign_ * distance_;
        }
    };

    double get_t_f(){
        return T_;
    };
    // 同步各个自由度的运动时间
    void set_t_f(double t_f_max){
        T_ = std::max(t_f_max, T_);
    };

  private:
    double T_;      // 半周期
    double distance_;  // 运动距离
    double sign_;   // 运动方向正负
};


template<typename TrajFunc>
class LineMotion {
  public:
    // 初始化: 速度, 方向, 距离
    LineMotion(double speed_factor, double distance, const std::array<double, 3> direction)
        :distance_(distance){
        // 方向向量归一化
        double norm = std::sqrt(std::pow(direction[0], 2) + std::pow(direction[1], 2) \
            + std::pow(direction[2], 2));
        if (norm == 0) {
            throw std::invalid_argument("Invalid direction.");
        }
        direction_[0] = direction[0] / norm;
        direction_[1] = direction[1] / norm;
        direction_[2] = direction[2] / norm;

        // 创建轨迹函数
        vel_m_ *= speed_factor;
        acc_m_ *= speed_factor;
        traj_fun = TrajFunc(vel_m_, acc_m_, distance_);
    }

    franka::CartesianPose operator()(const franka::RobotState& robot_state, franka::Duration period){
        // 读取初始状态
        time_ += period.toSec();  // 运动过程经历时长 t
        if (time_ == 0.0) {
            // O_T_EE_c: 运动轨迹发生器的上一条指令末端执行器相对于基座的位姿.
            init_pose_ = robot_state.O_T_EE_c;
        }

        // 计算 t 时刻状态
        double delta_dist = traj_fun(time_);
        std::array<double, 16> new_pose = init_pose_;
        new_pose[12] += delta_dist * direction_[0];
        new_pose[13] += delta_dist * direction_[1];
        new_pose[14] += delta_dist * direction_[2];

        bool motion_finished = (delta_dist == distance_);
        franka::CartesianPose output(new_pose);
        output.motion_finished = motion_finished;
        return output;
    }

  private:
    double time_ = 0.0;
    double distance_ = 0.0;
    double vel_m_ = 0.02;
    double acc_m_ = 0.1;
    std::array<double, 3> direction_;
    std::array<double, 16> init_pose_;
    TrajFunc traj_fun;
};


/**
* JointMotion
*/
template<typename TrajFunc>
class JointMotion{
  public:
    JointMotion(double speed_factor, const std::array<double, 7> q_goal)
        : q_goal_(q_goal.data()) {
        dq_max_ *= speed_factor;
        ddq_max_ *= speed_factor;
    };

    franka::JointPositions operator()(const franka::RobotState& robot_state,
        franka::Duration period) {
        time_ += period.toSec();  // 运动过程经历时长 t

        if (time_ == 0.0) {
            q_start_ = Vector7d(robot_state.q_d.data());  // 上一个时刻的 q_d 作为关节初始状态
            q_distance_ = q_goal_ - q_start_;
            synchronize_();
        }

        std::array<double, 7> q_output;
        Eigen::VectorXd::Map(&q_output[0], 7) = q_start_;
        // 计算 t 时刻状态
        bool motion_finished = true;
        for(size_t i=0; i<7; ++i){
            if(!q_finish_[i]){
                double delta_q = traj_funs[i](time_);
                q_output[i] += delta_q;
                q_finish_[i] = (delta_q == q_distance_[i]);
                motion_finished &= q_finish_[i];
            }
        }

        franka::JointPositions output(q_output);
        output.motion_finished = motion_finished;
        return output;
    }

  private:
    using Vector7d = Eigen::Matrix<double, 7, 1, Eigen::ColMajor>;

    void synchronize_(){
        double t_f_max = 0.0, t_f;
        for(size_t i=0; i<7; ++i){
            if(std::abs(q_distance_[i]) > FinishThreshold) {
                traj_funs[i] = TrajFunc(dq_max_[i], ddq_max_[i], q_distance_[i]);
                t_f = traj_funs[i].get_t_f();
                if(t_f > t_f_max) t_f_max = t_f;
            } else{
                q_finish_[i] = true;
            }
        }
        for(size_t i=0; i<7; ++i){
            if(!q_finish_[i]){
                traj_funs[i].set_t_f(t_f_max);
            }
        }
    }

    Vector7d dq_max_ = (Vector7d() << 2.0, 2.0, 2.0, 2.0, 2.5, 2.5, 2.5).finished(); // 各个关节最大速度
    Vector7d ddq_max_ = (Vector7d() << 5, 5, 5, 5, 5, 5, 5).finished();  // 加速度阶段最大加速度
    Vector7d q_start_;
    Vector7d q_goal_;
    Vector7d q_distance_;
    std::array<bool, 7> q_finish_{};
    std::array<TrajFunc, 7> traj_funs;
    double time_ = 0.0;
    static constexpr double FinishThreshold = 1e-6;  // 判断发生运动的阈值
};


